import urllib.request  # 导入 urllib.request 模块，用于进行网络请求
import zipfile  # 导入 zipfile 模块，用于处理 ZIP 压缩文件
import os  # 导入 os 模块，用于操作系统相关的操作
from pathlib import Path  # 导入 Path 类，用于处理文件路径
import pandas as pd  # 导入 pandas 库，并使用别名 pd

# 设置要下载的 ZIP 文件的 URL
url = "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"
# 设置 ZIP 文件的保存路径
zip_path = "sms_spam_collection.zip"
# 设置解压后的文件路径
extracted_path = "sms_spam_collection"
# 设置最终数据文件的路径
data_file_path = Path(extracted_path) / "SMSSpamCollection.tsv"

# 定义一个函数用于下载和解压垃圾短信数据
def download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path):
    # 如果数据文件已经存在
    if data_file_path.exists():
        # 打印提示信息并跳过下载和解压步骤
        print(f"{data_file_path} already exists. Skipping download and extraction.")
        return

    # 使用 urllib.request.urlopen 打开 URL 进行网络请求
    with urllib.request.urlopen(url) as response:
        # 以二进制写入模式打开 ZIP 文件
        with open(zip_path, "wb") as out_file:
            # 将网络响应的内容写入 ZIP 文件
            out_file.write(response.read())
            # 这一步完成了文件的下载

    # 使用 zipfile.ZipFile 打开 ZIP 文件
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        # 将 ZIP 文件解压到指定的路径
        zip_ref.extractall(extracted_path)

    # 获取原始文件路径
    original_file_path = Path(extracted_path) / "SMsSpamCollection"
    # 将原始文件重命名为带有.tsv 扩展名的文件
    os.rename(original_file_path, data_file_path)

# 调用函数进行下载和解压操作
download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)

# 使用 pandas 的 read_csv 函数读取数据文件
df = pd.read_csv(data_file_path, sep="\t", header=None, names=["Label", "Text"])
print(df)

def create_balanced_dataset(df):
    # 计算标签为"spam"的实例数量
    num_spam = df[df["Label"] == "spam"].shape[0]
    # 从标签为"ham"的实例中随机采样，数量与"spam"的实例数量相同，设置随机种子为 123
    ham_subset = df[df["Label"] == "ham"].sample(num_spam, random_state=123)
    # 将标签为"spam"的子集和随机采样后的"ham"子集合并
    balanced_df = pd.concat([df[df["Label"] == "spam"], ham_subset])
    # 打印平衡后数据集中的标签数量分布
    print(balanced_df["Label"].value_counts())
    # 返回平衡后的数据集
    return balanced_df

# 使用上述函数创建平衡数据集
balanced_df = create_balanced_dataset(df)
print(balanced_df["Label"].value_counts())

def random_split(df, train_frac, validation_frac):
    # 随机打乱整个 DataFrame，设置随机种子为 123
    df = df.sample(frac=1, random_state=123).reset_index(drop=True)
    # 计算验证集的结束索引位置，通过数据集长度乘以验证集比例得到验证集长度，再加上训练集结束索引位置
    # 根据比例确定训练集结束索引位置
    train_end = int(len(df) * train_frac)
    # 按照索引位置分割 DataFrame，得到训练集
    train_df = df[:train_end]
    validation_end = train_end + int(len(df) * validation_frac)
    # 按照索引位置分割 DataFrame，得到验证集
    validation_df = df[train_end:validation_end]
    # 按照索引位置分割 DataFrame，得到测试集（剩余部分）
    test_df = df[validation_end:]
    # 返回训练集、验证集和测试集
    return train_df, validation_df, test_df

# 使用 random_split 函数对平衡后的数据集进行分割，训练集比例为 0.7，验证集比例为 0.1，测试集比例默认为剩余的 0.2
train_df, validation_df, test_df = random_split(balanced_df, 0.7, 0.1)

train_df.to_csv("train.csv", index=None)
validation_df.to_csv("validation.csv", index=None)
test_df.to_csv("test.csv", index=None)


